Feat: add dictionaries as a supported group column type#23187
Feat: add dictionaries as a supported group column type#23187Rich-T-kid wants to merge 25 commits into
Conversation
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
2b60132 to
1ffc7f0
Compare
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
3f7ff57 to
e6b6dce
Compare
|
@kumarUjjawal could you run the dictionary benchmarks on this PR? Thx |
| } | ||
| } | ||
| DataType::Dictionary(key_dt, value_dt) => { | ||
| let new_field = Field::new("", *value_dt.clone(), true); |
There was a problem hiding this comment.
Since this field is never read again it may be fine to ignore the name field.
should be weary of similar issues to #21765 (comment)
There was a problem hiding this comment.
Kind of annoying that make_group_column takes a field instead of a DataType. Maybe we can change that in a follow up PR?
|
@kumarUjjawal wanted to bump this 😄 |
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (eb41915) to 01bf68c (merge-base) diff using: dictionary_group_values File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagedictionary_group_values — base (merge-base)
dictionary_group_values — branch
File an issue against this benchmark runner |
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
eb41915 to
770abfe
Compare
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
770abfe to
243a557
Compare
Rich-T-kid
changed the title
|
@codex review |
@Rich-T-kid Thank you! I have been sick so I won't be available for review. I will probably get back next week. |
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
d4ea0f9 to
894a0a6
Compare
| pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> { | ||
| inner: Box<dyn GroupColumn>, | ||
| null_array: ArrayRef, | ||
| // Mutex is required because `vectorized_equal_to` takes `&self` (GroupColumn trait |
There was a problem hiding this comment.
RefCell doesn't work here because its not sync
There was a problem hiding this comment.
std::sync::RwLock may also be an option
|
PR is ready for review 🚀 |
Did you run them yourself locally? What are your results? |
@geoffreyclaude yes. These are the results compared to main |
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (894a0a6) to 01bf68c (merge-base) diff using: dictionary_group_values File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagedictionary_group_values — base (merge-base)
dictionary_group_values — branch
File an issue against this benchmark runner |
| let unique_lhs_row_count = | ||
| lhs_rows.iter().collect::<hashbrown::HashSet<_>>().len(); | ||
| if unique_lhs_row_count <= rhs_rows.len() / CACHE_USE_THRESHOLD { | ||
| let mut comparison_cache = self.key_to_group_cache.lock().unwrap(); | ||
| for (position, (&lhs_row, &rhs_row)) in | ||
| lhs_rows.iter().zip(rhs_rows).enumerate() | ||
| { | ||
| if !equal_to_results.get_bit(position) { | ||
| continue; | ||
| } | ||
| let dict_key = dict.key(rhs_row); | ||
| let is_equal = *comparison_cache | ||
| .entry((lhs_row, dict_key)) | ||
| .or_insert_with(|| match dict_key { | ||
| None => self.inner.equal_to(lhs_row, &self.null_array, 0), | ||
| Some(key) => self.inner.equal_to(lhs_row, dict_values, key), | ||
| }); |
There was a problem hiding this comment.
This is causing more overhead than its saving.
There was a problem hiding this comment.
hashset insertions + mutext lock + hashsing two usize's
O(size * hash) O(1) O(~ unqiue values)
There was a problem hiding this comment.
I think the best thing to do here to to avoid any regressions. That means making this implementation a thin wrapper over an inner GroupColumns implementation. a follow up PR can try to optimize the low-cardinality case similar to #21765
Rich-T-kid
changed the title
| } | ||
| } | ||
| DataType::Dictionary(key_dt, value_dt) => { | ||
| let new_field = Field::new("", *value_dt.clone(), true); |
There was a problem hiding this comment.
Kind of annoying that make_group_column takes a field instead of a DataType. Maybe we can change that in a follow up PR?
| // `get_combined` appends the null sentinel as the final element, so any | ||
| // null dictionary key maps to that last index. | ||
| let combined_with_null_sentinel = | ||
| concat(&[dict.values(), self.null_array.as_ref()]) |
There was a problem hiding this comment.
The concat isn't great tbh bc it reallocates the whole array, but I don't see how we can avoid it. Since the benchmarks are good, it's probably fine.
It's unfortunate that keys might be null. Ideally the dict arrives "normalized" where all keys are valid and may point to a null value.
How likely is it that there's null keys? Would it help to see how often .unwrap_or(null_sentinel_index) happens? If it happens 0 times, we can avoid concat?
There was a problem hiding this comment.
The concat isn't great tbh bc it reallocates the whole array, but I don't see how we can avoid it. Since the benchmarks are good, it's probably fine.
I agree, for a 8192 size batch of keys with just a single null value it will cause the entire 8192 batch to be copied to a new array just to append 1 null at the end. Invectorized_append() we can avoid copying value arrays blindly by doing a Arc::ptr_eq() check. this is a cheap comparison that can save an O(target_batch_size) + 1 copy for every call. as for vectorized_equal_to() theres not much we can do here since it takes a &self meaning we wouldn't be able to re-cache the values array.
It's unfortunate that keys might be null. Ideally the dict arrives "normalized" where all keys are valid and may point to a null value.
How likely is it that there's null keys? Would it help to see how often.unwrap_or(null_sentinel_index)happens? If it happens 0 times, we can avoidconcat?
This is what I tried avoiding by adding the null branch case. Fundemantally we have to account for the possibility of keys being null.
arrows format doc:
The memory layout for a dictionary-encoded array is the same as that of a primitive integer layout.
the primitive array section states:
A primitive value array represents an array of values each having the same physical slot width typically measured in bytes, though the spec also provides for bit-packed types (e.g. boolean values encoded in bits).
Internally, the array contains a contiguous memory buffer whose total size is at least as large as the slot width multiplied by the array length. For bit-packed types, the size is rounded up to the nearest byte.
The associated validity bitmap is contiguously allocated (as described above) but does not need to be adjacent in memory to the values buffer.
Example Layout: Int32 Array
For example a primitive array of int32s:
[1, null, 2, 4, 8]
Unlike a primitive, where nulls live in a single validity bitmap, a dictionary array's logical nulls can come from two independent places: the keys array's validity bitmap, or a null entry in the values (dictionary) array that a key points to. The field's nullable flag doesn't tell us which it's only an advisory "may contain nulls," same as for any type. So to know the actual null situation we have to inspect both the keys and the values arrays.
@jayshrivastava
|
Rich-T-kid
force-pushed
the
rich-T-kid/dictionary-groupValuesColumn-impl
branch
from
1f4816e to
e2e79ba
Compare
|
@jayshrivastava I pushed a revision could you take another look? |
|
@Rich-T-kid Your optimization relies on the pointer being the same, generally, what are the odds of that? I suggested earlier to count how many times One way to find out is to run the benches and print out how many times each of those things happen, assuming there's enough interesting cases in the benches. Pointer checking is a bit of a smell. We used to do it in dynamic filtering to check if two filters are the same but everyone disliked that. We have since moved away from pointer equality checking. |
|
Also, if there is already a null value in the values array, then we don't need to append a sentinel. Avoiding concat if there's no null keys or if there's already a null value might be good enough. It could avoid concating more than the pointer check. |
This make sense. updating the PR 👍 |
in the benchmarks the null rate is set to 0% and 10%. in the case where there are no nulls in the current batch there is no concat operation. the alternative to using concat is to use a non-vectorized approach. // Null keys but values has no null entry: scalar fallback per row.
for &row_index in rows {
match dict.key(row_index) {
None => self.inner.append_val(&self.null_array, 0)?,
Some(key) => self.inner.append_val(dict.values(), key)?,
}
}
make sense, Ive removed the pointer check, as well as the concat call. |
|
I believe this PR is ready for a maintainer to take a look 👍 |
|
run benchmark dictionary_group_values |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing rich-T-kid/dictionary-groupValuesColumn-impl (fb7856f) to 01bf68c (merge-base) diff using: dictionary_group_values File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagedictionary_group_values — base (merge-base)
dictionary_group_values — branch
File an issue against this benchmark runner |
| } | ||
| } | ||
|
|
||
| fn into_dict(values: ArrayRef) -> ArrayRef { |
There was a problem hiding this comment.
The emitted DictionaryArray uses identity keys (key[i] = i), so output is fully expanded with no value sharing. Could you add a short doc-comment explaining this?
There was a problem hiding this comment.
valid_bounds is guarded with assert! which panics. A Dictionary(Int8, Utf8) column with >127 distinct groups would cause a panic in production rather than a graceful error. All other error paths here use Result. converting into_dict to -> Result and returning internal_err! on overflow, or checking bounds at intern time so build / take_n stay infallible.
| &value_indices, | ||
| equal_to_results, | ||
| ); | ||
| } else if let Some(sentinel) = Self::null_sentinel_index(dict_values) { |
There was a problem hiding this comment.
a null key in rhs_rows is mapped to sentinel_idx inside dict_values (the batch's values array), and then passed to self.inner.vectorized_equal_to . But the null that was originally appended went through self.inner.append_val(&self.null_array, 0), a null from a different array than this batch's dict_values .
Can you confirm that all inner builder implementations (bytes, primitive, etc.) short-circuit on null equality before reading the array content so comparing "null from null_array " vs "null from dict_values [sentinel]" always returns true ?
| Arc::new(DictionaryArray::<K>::new(keys, values)) | ||
| } | ||
|
|
||
| /// Returns the index of the first null in `values`, or `None` if there are no nulls. |
There was a problem hiding this comment.
can we use use bit-level null iteration?
5 participants
Which issue does this PR close?
Rationale for this change
This PR introduces a specialized
GroupColumnimplementation for dictionary-typed columns insideGroupValuesColumn, allowing dictionary columns to participate in the columnar, vectorized aggregation path instead of the row-based fallback.The Implementation is only about 175+ lines of code. the remaining LOC is adding extensive test at the
GroupColumntrait level as well as testing theGroupValuesColumnGroupValues trait and how it inter-opts with multi-dictionary group by's.What changes are included in this PR?
DictionaryGroupValueBuilderstruct implementing theGroupColumntrait forDictionary-typed group-by columns, supporting a configurable subset of value typesGroupValuesColumn::try_new(thematches!block) to acceptDictionary(_, value_type)wherevalue_typeis already supported.emitAre these changes tested?
yes. a majority of this PR is test
Are there any user-facing changes?
no. this is a pure perf boost for users.